from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.core.llms import ChatMessage

# 定义模型和分词器的路径
# model_path = r"D:\AIProject\modelscope\Qwen\Qwen2___5-0___5B-Instruct"
# tokenizer_path = r"D:\AIProject\modelscope\Qwen\Qwen2___5-0___5B-Instruct"
model_path = "/Users/wufagang/PycharmProjects/llamaindex-rag/.t/modelscope/Qwen/Qwen2___5-0___5B-Instruct"
tokenizer_path = "/Users/wufagang/PycharmProjects/llamaindex-rag/modelscope/Shanghai_AI_Laboratory/internlm2-chat-1_8b"
#  使用 HuggingFace 加载本地大模型
# 参数
# 第一个：指定要加载的模型的路径或名称
# 第二个：指定要加载的分词器的路径或名称
# 第三个： 传递给模型加载函数的额外参数  trust_remote_code 是否信任远程代码
# 第四个： 传递给分词器加载函数的额外参数
llm = HuggingFaceLLM(
    # 给定的是本地模型的路径
    model_name=model_path,
    tokenizer_name=tokenizer_path,
    model_kwargs={"trust_remote_code":True,"torch_dtype": "auto","load_in_8bit": True},
    tokenizer_kwargs={"trust_remote_code":True}
)

# 调用模型
try:
    rsp = llm.chat(messages=[ChatMessage(content="xtuner是什么？")])
    print(rsp)
except Exception as e:
    print(f"Error: {e}")